import numpy as np
from scipy.stats import chi2, ncx2

# manual_seed =42
# 1
lambda_matrix=0.5
T_initial=15

def find_best_intervention(graph, inverse_m, hat_theta_x, hat_theta_y, norm_matrix_mx, rho):
    max_y = - 9999
    for i in range(graph.num_parents_y):
        hat_theta_x[i] = hat_theta_x[i] + rho / norm_matrix_mx[i] if norm_matrix_mx[i] > 0 else 1
    best_intervention = []
    for j in range(np.power(graph.num_parents_y, graph.k)):
        intervened_indexes = []
        for k in range(graph.k):
            index = int(j / np.power(graph.num_parents_y, k)) % graph.num_parents_y
            if len(intervened_indexes) == 0 or index > intervened_indexes[-1]:
                intervened_indexes.append(index)
        if len(intervened_indexes) < graph.k:
            continue
        tilde_theta_x = hat_theta_x.copy()
        for i in intervened_indexes:
            tilde_theta_x[int(i), 0] = 1
        product_result=np.matmul(tilde_theta_x.T, hat_theta_y)
        # print(" product_result",product_result)
        # print(" hat_theta_y", hat_theta_y)
        expected_y = rho * np.sqrt(np.matmul(np.matmul(tilde_theta_x.T, inverse_m), tilde_theta_x)) + np.matmul(
            tilde_theta_x.T, hat_theta_y)[0, 0, 0]
        if expected_y > max_y:
            max_y = expected_y
            best_intervention = intervened_indexes
  
    return best_intervention



def find_best_intervention_cluster_y(graph, inverse_m_one,inverse_m_two, hat_theta_x, hat_theta_y_one,hat_theta_y_two, norm_matrix_mx, rho):
    max_y = - 9999
    for i in range(graph.num_parents_y):
        hat_theta_x[i] = hat_theta_x[i] + rho / norm_matrix_mx[i] if norm_matrix_mx[i] > 0 else 1
    best_intervention = []
    for j in range(np.power(graph.num_parents_y, graph.k)):
        intervened_indexes = []
        for k in range(graph.k):
            index = int(j / np.power(graph.num_parents_y, k)) % graph.num_parents_y
            if len(intervened_indexes) == 0 or index > intervened_indexes[-1]:
                intervened_indexes.append(index)
        if len(intervened_indexes) < graph.k:
            continue
        tilde_theta_x = hat_theta_x.copy()
        for i in intervened_indexes:
            tilde_theta_x[int(i), 0] = 1
     
        tilde_theta_x_one=tilde_theta_x[:3]
        tilde_theta_x_two=tilde_theta_x[-3:]
      
        expected_y_one = rho * np.sqrt(np.matmul(np.matmul(  tilde_theta_x_one.T, inverse_m_one), tilde_theta_x_one)) + np.matmul(
            tilde_theta_x_one.T, hat_theta_y_one)[0, 0, 0]
        expected_y_two = rho * np.sqrt(np.matmul(np.matmul(  tilde_theta_x_two.T, inverse_m_two), tilde_theta_x_two)) + np.matmul(
            tilde_theta_x_two.T, hat_theta_y_two)[0, 0, 0]
        expected_y= expected_y_one+expected_y_two
        if expected_y > max_y:
            max_y = expected_y
            best_intervention = intervened_indexes
    # print(" best_intervention:", best_intervention)
    return best_intervention



def find_best_intervention_cluster_two(graph, inverse_m_one,inverse_m_two,inverse_m_three, hat_theta_x, hat_theta_y_one,hat_theta_y_two, hat_theta_y_three,norm_matrix_mx, rho):
    max_y = - 9999
    for i in range(graph.num_parents_y):
        hat_theta_x[i] = hat_theta_x[i] + rho / norm_matrix_mx[i] if norm_matrix_mx[i] > 0 else 1
    best_intervention = []
    for j in range(np.power(graph.num_parents_y, graph.k)):
        intervened_indexes = []
        for k in range(graph.k):
            index = int(j / np.power(graph.num_parents_y, k)) % graph.num_parents_y
            if len(intervened_indexes) == 0 or index > intervened_indexes[-1]:
                intervened_indexes.append(index)
        if len(intervened_indexes) < graph.k:
            continue
        tilde_theta_x = hat_theta_x.copy()
        for i in intervened_indexes:
            tilde_theta_x[int(i), 0] = 1
     
        tilde_theta_x_one=tilde_theta_x[:2]
        tilde_theta_x_two=tilde_theta_x[-2:]
        tilde_theta_x_three=tilde_theta_x[2:4]
      
        expected_y_one = rho * np.sqrt(np.matmul(np.matmul(  tilde_theta_x_one.T, inverse_m_one), tilde_theta_x_one)) + np.matmul(
            tilde_theta_x_one.T, hat_theta_y_one)[0, 0, 0]
        expected_y_two = rho * np.sqrt(np.matmul(np.matmul(  tilde_theta_x_two.T, inverse_m_two), tilde_theta_x_two)) + np.matmul(
            tilde_theta_x_two.T, hat_theta_y_two)[0, 0, 0]

        expected_y_three = rho * np.sqrt(np.matmul(np.matmul(  tilde_theta_x_three.T, inverse_m_three), tilde_theta_x_three)) + np.matmul(
            tilde_theta_x_three.T, hat_theta_y_three)[0, 0, 0]
        expected_y= expected_y_one+expected_y_two+expected_y_three
        if expected_y > max_y:
            max_y = expected_y
            best_intervention = intervened_indexes
    # print(" best_intervention:", best_intervention)
    return best_intervention



def find_best_intervention_cluster_component(graph, inverse_m_list, hat_theta_x, hat_theta_list, norm_matrix_mx, rho):
    max_y = - 9999
    for i in range(graph.num_parents_y):
        hat_theta_x[i] = hat_theta_x[i] + rho / norm_matrix_mx[i] if norm_matrix_mx[i] > 0 else 1
    best_intervention = []
    for j in range(np.power(graph.num_parents_y, graph.k)):
        intervened_indexes = []
        for k in range(graph.k):
            index = int(j / np.power(graph.num_parents_y, k)) % graph.num_parents_y
            if len(intervened_indexes) == 0 or index > intervened_indexes[-1]:
                intervened_indexes.append(index)
        if len(intervened_indexes) < graph.k:
            continue
        tilde_theta_x = hat_theta_x.copy()
        for i in intervened_indexes:
            tilde_theta_x[int(i), 0] = 1

        expected_y=0
        for j in range(6):
     
            expected_y_now = rho * np.sqrt(np.matmul(np.matmul(  tilde_theta_x[j].T, inverse_m_list[j]), tilde_theta_x[j])) + tilde_theta_x[j][0]*hat_theta_list[j][0, 0, 0]

        if expected_y > max_y:
            max_y = expected_y
            best_intervention = intervened_indexes
    # print(" best_intervention:", best_intervention)
    return best_intervention



def pair_oracle_3(hat_theta, inverse_m, rho, parents):
    # print(parents, inverse_m)
    # print(np.matmul(parents.T, inverse_m))
    return (rho * np.sqrt(np.matmul(np.matmul(parents.T, inverse_m), parents)[0, 0]) + np.matmul(
        parents.T, hat_theta))[0, 0, 0]


def find_best_intervention_two_layer(graph, inverse_m1, inverse_m2, inverse_m3, inverse_my, hat_theta_1, hat_theta_2,
                                     hat_theta_3, hat_theta_4, hat_theta_y, norm_matrix_mx, rho):
    max_y = - 9999
    for i in range(2):
        hat_theta_1[i] = hat_theta_1[i] + rho / norm_matrix_mx[i] if norm_matrix_mx[i] > 0 else 1
    best_intervention = []
    for j in range(np.power(5, graph.k)):
        intervened_indexes = []
        for k in range(graph.k):
            index = int(j / np.power(5, k)) % 5
            if len(intervened_indexes) == 0 or index > intervened_indexes[-1]:
                intervened_indexes.append(index)
        if len(intervened_indexes) < graph.k:
            continue
        x1 = hat_theta_1[0, 0] if 0 not in intervened_indexes else 1
        x2 = hat_theta_1[1, 0] if 1 not in intervened_indexes else 1
        x3 = pair_oracle_3(hat_theta_2, inverse_m1, rho,
                           np.array([[1, x1, x2]]).T) if 2 not in intervened_indexes else 1
        x4 = pair_oracle_3(hat_theta_3, inverse_m2, rho,
                           np.array([[1, x1, x2]]).T) if 3 not in intervened_indexes else 1
        x5 = pair_oracle_3(hat_theta_4, inverse_m3, rho,
                           np.array([[1, x1, x2]]).T) if 4 not in intervened_indexes else 1
        expected_y = pair_oracle_3(hat_theta_y, inverse_my, rho, np.array([[x3, x4, x5]]).T)
        # print(x1, x2, x3,x4,x5,rho, intervened_indexes, expected_y, graph.expect_y(intervened_indexes))
        if expected_y > max_y:
            max_y = expected_y
            best_intervention = intervened_indexes
    # print(best_intervention)
    return best_intervention



def run_blm_lr_partial(graph_list):
    graph=graph_list[0]
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}

    if graph.type == 'parallel':
        for m in range(len(graph_list)):
            matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
            for i in range(graph.num_parents_y):
                matrix_m[i][i] = 1
           
            matrix_m_list[str(m)]=  matrix_m
            matrix_mx = np.zeros(graph.num_parents_y)
            matrix_mx_list[str(m)]=  matrix_mx
      
        
            intervened_times = np.zeros(graph.num_parents_y)
            intervened_times_list[str(m)]=   intervened_times
            by = np.array([np.zeros(graph.num_parents_y)]).T
            by_list[str(m)]= by
            payoff_list = [0]
            total_expected_payoff = 0
        
        # we have a graph list
        #
        for i in range(graph.T):
            for m in range(len(graph_list)):
                graph=graph_list[m]
                matrix_m=  matrix_m_list[str(m)]
                matrix_mx=  matrix_mx_list[str(m)]

                intervened_times=intervened_times_list[str(m)]
                by= by_list[str(m)] 



                inverse_m = np.linalg.inv(matrix_m)
                hat_theta_x = np.zeros(graph.num_parents_y)
                for j in range(graph.num_parents_y):
                    hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
                hat_theta_x = np.array([hat_theta_x]).T
                norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])
                hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
                rho = graph.compute_rho_lr(i + 1)
                best_intervention = find_best_intervention(graph, inverse_m, hat_theta_x, hat_theta_y, norm_matrix_mx, rho)
                vy, y = graph.simulate(best_intervention)
                total_expected_payoff += graph.expect_y(best_intervention)
                
               
               
                matrix_m += np.matmul(vy, vy.T)
                by += y * vy
                for j in best_intervention:
                    intervened_times[int(j)] += 1
                for j in range(graph.num_parents_y):
                    if j not in best_intervention:
                        matrix_mx[j] += vy[j]

                matrix_m_list[str(m)]= matrix_m
                matrix_mx_list[str(m)]=matrix_mx
                intervened_times_list[str(m)]=intervened_times
                by_list[str(m)] =by

            for m in range(len(graph_list)):
                if (m==0):
                    matrix_mx_list[str(m)]=  matrix_mx_list[str(m)]+ matrix_mx_list[str(1)]
                    intervened_times_list[str(m)]= intervened_times_list[str(m)]+intervened_times_list[str(1)]
                    matrix_m_list[str(m)]=   matrix_m_list[str(m)]+   matrix_m_list[str(1)]
                    by_list[str(m)] =  by_list[str(m)]+  by_list[str(1)]
                    
                        

                if (m==1):
                    matrix_mx_list[str(m)]=  matrix_mx_list[str(0)]
                    intervened_times_list[str(m)]= intervened_times_list[str(0)]
                    matrix_m_list[str(m)]=matrix_m_list[str(0)]
                    by_list[str(m)] =   by_list[str(0)] 
                    
                   

                if (m==2):
                    matrix_mx_list[str(m)]=  matrix_mx_list[str(m)]+ matrix_mx_list[str(3)]
                    intervened_times_list[str(m)]= intervened_times_list[str(m)]+intervened_times_list[str(3)]
                    matrix_m_list[str(m)]=   matrix_m_list[str(m)]+   matrix_m_list[str(3)]
                    by_list[str(m)] =  by_list[str(m)]+  by_list[str(3)]
                    
                    
                if (m==3):
                    matrix_mx_list[str(m)]=  matrix_mx_list[str(2)]
                    intervened_times_list[str(m)]= intervened_times_list[str(2)]
                    matrix_m_list[str(m)]=matrix_m_list[str(2)]
                    by_list[str(m)] =   by_list[str(2)] 
                  
                       

            if i % 1000 == 99:
                print("partial")
              
                print("round:",i)

                print("reward:", total_expected_payoff)
                payoff_list.append(total_expected_payoff)

                
            # print(best_intervention)
        # print(np.matmul(np.linalg.inv(matrix_m), np.array([by]).T))
        return payoff_list, total_expected_payoff


def run_blm_lr_partition_Fed(graph_list,best_y_list):
    graph=graph_list[0]
    matrix_m_one_list={}
    matrix_m_two_list={}
    by_one_list={}
    by_two_list={}
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}

    if graph.type == 'parallel':
        for m in range(len(graph_list)):
            matrix_m_one = np.zeros((3, 3))
            matrix_m_two = np.zeros((3, 3))
            
          
            for i in range(3):
                matrix_m_one[i][i] = 1
                matrix_m_two[i][i] = 1
            by_one = np.array([np.zeros(3)]).T
            by_two = np.array([np.zeros(3)]).T

            matrix_m_one_sum= np.zeros((3, 3))
            matrix_m_two_sum= np.zeros((3, 3))
            for i in range(3):
                matrix_m_one_sum[i][i] = 1
                matrix_m_two_sum[i][i] = 1
            by_one_sum = np.array([np.zeros(3)]).T
            by_two_sum = np.array([np.zeros(3)]).T

            matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
            for i in range(graph.num_parents_y):
                matrix_m[i][i] = 1

            
        
            matrix_m_one_list[str(m)]=  matrix_m_one 
            matrix_m_two_list[str(m)]=  matrix_m_two
            by_one_list[str(m)]= by_one
            by_two_list[str(m)]= by_two

            matrix_m_list[str(m)]=  matrix_m
            matrix_mx = np.zeros(graph.num_parents_y)
            matrix_mx_list[str(m)]=  matrix_mx
      
        
            intervened_times = np.zeros(graph.num_parents_y)
            intervened_times_list[str(m)]=   intervened_times
            by = np.array([np.zeros(graph.num_parents_y)]).T
            by_list[str(m)]= by
            payoff_list = [0]
            total_expected_payoff = 0
            total_optimal_reward=0
        
        # we have a graph list
        #
        T_remaining=graph.T-T_initial
        for i in range(T_initial):
            for m in range(len(graph_list)):
                graph=graph_list[m]
                best_intervention=[0,1,2]
                total_expected_payoff += graph.expect_y(best_intervention)
                # total_expected_payoff+=0
                total_optimal_reward+=best_y_list[m]

    
        for i in range(T_remaining):
            for m in range(len(graph_list)):
                graph=graph_list[m]
                # matrix_m=  matrix_m_list[str(m)]

                matrix_mx=  matrix_mx_list[str(m)]

                matrix_m_one=matrix_m_one_list[str(m)]
                matrix_m_two=matrix_m_two_list[str(m)]
                by_one= by_one_list[str(m)]
                by_two= by_two_list[str(m)]
               

                intervened_times=intervened_times_list[str(m)]
            

                # inverse_m = np.linalg.inv(matrix_m)
                # print("matrix_m:",  matrix_m.shape)
                # inverse_m_one=np.linalg.inv(matrix_m_one)
               
                inverse_m_one = np.linalg.inv(matrix_m_one)
                inverse_m_two = np.linalg.inv(matrix_m_two)



                hat_theta_x = np.zeros(graph.num_parents_y)
                for j in range(graph.num_parents_y):
                    hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
                hat_theta_x = np.array([hat_theta_x]).T
                norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])

            
             
                hat_theta_y_one = np.matmul(inverse_m_one, np.array([by_one]).T)
                hat_theta_y_two = np.matmul(inverse_m_two, np.array([by_two]).T)
              
                rho = graph.compute_rho_lr(i + 1)

                best_intervention = find_best_intervention_cluster_y(graph,inverse_m_one,inverse_m_two, hat_theta_x,hat_theta_y_one,hat_theta_y_two,norm_matrix_mx, rho)
                
                vy, y,y_g_one,y_g_two= graph.simulate_cluster_y(best_intervention)
          
                vy_g_one=vy[:3]
                vy_g_two=vy[-3:]
               

                total_expected_payoff += graph.expect_y(best_intervention)
                total_optimal_reward+=best_y_list[m]

                matrix_m_one_sum+= np.matmul(vy_g_one,vy_g_one.T)
                matrix_m_two_sum+= np.matmul(vy_g_two,vy_g_two.T)
                by_one_sum+=y_g_one*vy_g_one
                by_two_sum+=y_g_two*vy_g_two
               
                # matrix_m += np.matmul(vy, vy.T)
                # by += y * vy
                for j in best_intervention:
                    intervened_times[int(j)] += 1
                for j in range(graph.num_parents_y):
                    if j not in best_intervention:
                        matrix_mx[j] += vy[j]

              

                matrix_mx_list[str(m)]=matrix_mx
                intervened_times_list[str(m)]=intervened_times

              
            for m in range(len(graph_list)):
                matrix_m_one_list[str(m)]=  matrix_m_one_sum
                matrix_m_two_list[str(m)]=  matrix_m_two_sum
                by_one_list[str(m)]= by_one_sum
                by_two_list[str(m)]= by_two_sum




            if i % 100 == 99:
                print("partition_y_Fed")
              
                print("round:",i)

                print("reward:", total_expected_payoff)
                total_regret=total_optimal_reward-total_expected_payoff
                payoff_list.append(total_regret)

                
            # print(best_intervention)
        # print(np.matmul(np.linalg.inv(matrix_m), np.array([by]).T))
        return payoff_list, total_expected_payoff

def run_blm_lr_partition_cluster_vector(graph_list,best_y_list):
    graph=graph_list[0]
    matrix_m_one_list={}
    matrix_m_two_list={}
    by_one_list={}
    by_two_list={}
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}

    if graph.type == 'parallel':
        for m in range(len(graph_list)):
            matrix_m_one = np.zeros((3, 3))
            matrix_m_two = np.zeros((3, 3))
            
          
            for i in range(3):
                matrix_m_one[i][i] = lambda_matrix
                matrix_m_two[i][i] = lambda_matrix
            by_one = np.array([np.zeros(3)]).T
            by_two = np.array([np.zeros(3)]).T

            matrix_m_one_sum= np.zeros((3, 3))
            matrix_m_two_sum= np.zeros((3, 3))
            matrix_m_one_log= np.zeros((3, 3))
            matrix_m_two_log= np.zeros((3, 3))
            for i in range(3):
                matrix_m_one_sum[i][i] =lambda_matrix
                matrix_m_two_sum[i][i] = lambda_matrix
                matrix_m_one_log[i][i] = lambda_matrix
                matrix_m_two_log[i][i] = lambda_matrix
            by_one_sum = np.array([np.zeros(3)]).T
            by_two_sum = np.array([np.zeros(3)]).T
            by_one_log = np.array([np.zeros(3)]).T
            by_two_log = np.array([np.zeros(3)]).T

            matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
            for i in range(graph.num_parents_y):
                matrix_m[i][i] = 1

            
        
            matrix_m_one_list[str(m)]=  matrix_m_one 
            matrix_m_two_list[str(m)]=  matrix_m_two
            by_one_list[str(m)]= by_one
            by_two_list[str(m)]= by_two

            matrix_m_list[str(m)]=  matrix_m
            matrix_mx = np.zeros(graph.num_parents_y)
            matrix_mx_list[str(m)]=  matrix_mx
      
        
            intervened_times = np.zeros(graph.num_parents_y)
            intervened_times_list[str(m)]=   intervened_times
            by = np.array([np.zeros(graph.num_parents_y)]).T
            by_list[str(m)]= by
            payoff_list = [0]
            total_expected_payoff = 0
            total_optimal_reward=0
        
        # we have a graph list

        for i in range(graph.T):
            for m in range(len(graph_list)):
                graph=graph_list[m]
                # matrix_m=  matrix_m_list[str(m)]

                matrix_mx=  matrix_mx_list[str(m)]
                matrix_m_one=matrix_m_one_list[str(m)]
                matrix_m_two=matrix_m_two_list[str(m)]
                by_one= by_one_list[str(m)]
                by_two= by_two_list[str(m)]

                intervened_times=intervened_times_list[str(m)]
                by= by_list[str(m)] 

                # inverse_m = np.linalg.inv(matrix_m)
                # print("matrix_m:",  matrix_m.shape)
                # inverse_m_one=np.linalg.inv(matrix_m_one)
               
                inverse_m_one = np.linalg.inv(matrix_m_one)
                inverse_m_two = np.linalg.inv(matrix_m_two)



                hat_theta_x = np.zeros(graph.num_parents_y)
                for j in range(graph.num_parents_y):
                    hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
                hat_theta_x = np.array([hat_theta_x]).T
                norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])

            
             
                hat_theta_y_one = np.matmul(inverse_m_one, np.array([by_one]).T)
                hat_theta_y_two = np.matmul(inverse_m_two, np.array([by_two]).T)
              
                # hat_theta_y_con = np.concatenate((hat_theta_y_one,hat_theta_y_two), axis=1)
                
                # hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
                # print("hat_theta_y:",hat_theta_y)
                rho = graph.compute_rho_lr(i + 1)

                best_intervention = find_best_intervention_cluster_y(graph,inverse_m_one,inverse_m_two, hat_theta_x,hat_theta_y_one,hat_theta_y_two,norm_matrix_mx, rho)
                
                vy, y= graph.simulate(best_intervention)
          
                vy_g_one=vy[:3]
                vy_g_two=vy[-3:]

                y_g_one=y-np.dot(hat_theta_y_two.flatten(),vy_g_two.flatten())
                y_g_two=y-np.dot(hat_theta_y_one.flatten(),vy_g_one.flatten())
               

                
                total_expected_payoff += graph.expect_y(best_intervention)
                total_optimal_reward+=best_y_list[m]

                # matrix_m_one+= np.matmul(vy_g_one,vy_g_one.T)
                # matrix_m_two+= np.matmul(vy_g_two,vy_g_two.T)
                # by_one+=y_g_one*vy_g_one
                # by_two+=y_g_two*vy_g_two
                
                # matrix_m_one+= np.matmul(vy_g_one,vy_g_one.T)
                # matrix_m_two+= np.matmul(vy_g_two,vy_g_two.T)
                # by_one+=y_g_one*vy_g_one
                # by_two+=y_g_two*vy_g_two

                if (m%2==0):
                    matrix_m_one_sum+= np.matmul(vy_g_one,vy_g_one.T)
                    by_one_sum+=y_g_one*vy_g_one
                    matrix_m_two_sum+= np.matmul(vy_g_two,vy_g_two.T)
                    by_two_sum+=y_g_two*vy_g_two
                    # print("m_one",m)
                
                else:
                    matrix_m_one_log+= np.matmul(vy_g_one,vy_g_one.T)
                    by_one_log+=y_g_one*vy_g_one
                    matrix_m_two_log+= np.matmul(vy_g_two,vy_g_two.T)
                    by_two_log+=y_g_two*vy_g_two
                   
                    # print("m_two":m)
               
               
                # matrix_m += np.matmul(vy, vy.T)
                # by += y * vy
                for j in best_intervention:
                    intervened_times[int(j)] += 1
                for j in range(graph.num_parents_y):
                    if j not in best_intervention:
                        matrix_mx[j] += vy[j]

                matrix_m_one_list[str(m)]=  matrix_m_one 
                matrix_m_two_list[str(m)]=  matrix_m_two
                by_one_list[str(m)]= by_one
                by_two_list[str(m)]= by_two

                matrix_mx_list[str(m)]=matrix_mx
                intervened_times_list[str(m)]=intervened_times
            
            for m in range(len(graph_list)):
                if (m%2==0):
                    # print("m_one_agg:",m)
                    matrix_m_one_list[str(m)]= matrix_m_one_sum
                    by_one_list[str(m)]=by_one_sum
                    matrix_m_two_list[str(m)]= matrix_m_two_sum
                    by_two_list[str(m)]=by_two_sum

                else:
                    # print("m_one_agg:",m)
                    matrix_m_one_list[str(m)]= matrix_m_one_log
                    by_one_list[str(m)]=by_one_log
                    matrix_m_two_list[str(m)]= matrix_m_two_log
                    by_two_list[str(m)]=by_two_log
              
                


            if i % 100 == 99:
                print("partition_vec")
              
                print("round:",i)

                print("reward:", total_expected_payoff)
                total_regret=total_optimal_reward-total_expected_payoff
                payoff_list.append(total_regret)

                
            # print(best_intervention)
        # print(np.matmul(np.linalg.inv(matrix_m), np.array([by]).T))
        return payoff_list, total_expected_payoff

def run_blm_lr_partition_cluster_vector_explore(graph_list,best_y_list):
    graph=graph_list[0]
    matrix_m_one_list={}
    matrix_m_two_list={}
    by_one_list={}
    by_two_list={}
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}

   
    for m in range(len(graph_list)):
        matrix_m_one = np.zeros((3, 3))
        matrix_m_two = np.zeros((3, 3))
        
        
        for i in range(3):
            matrix_m_one[i][i] = lambda_matrix
            matrix_m_two[i][i] = lambda_matrix
        by_one = np.array([np.zeros(3)]).T
        by_two = np.array([np.zeros(3)]).T

        matrix_m_one_sum= np.zeros((3, 3))
        matrix_m_two_sum= np.zeros((3, 3))
        matrix_m_one_log= np.zeros((3, 3))
        matrix_m_two_log= np.zeros((3, 3))
        for i in range(3):
            matrix_m_one_sum[i][i] =lambda_matrix
            matrix_m_two_sum[i][i] = lambda_matrix
            matrix_m_one_log[i][i] = lambda_matrix
            matrix_m_two_log[i][i] = lambda_matrix
        by_one_sum = np.array([np.zeros(3)]).T
        by_two_sum = np.array([np.zeros(3)]).T
        by_one_log = np.array([np.zeros(3)]).T
        by_two_log = np.array([np.zeros(3)]).T

        matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
        for i in range(graph.num_parents_y):
            matrix_m[i][i] = 1

        
    
        matrix_m_one_list[str(m)]=  matrix_m_one 
        matrix_m_two_list[str(m)]=  matrix_m_two
        by_one_list[str(m)]= by_one
        by_two_list[str(m)]= by_two

        matrix_m_list[str(m)]=  matrix_m
        matrix_mx = np.zeros(graph.num_parents_y)
        matrix_mx_list[str(m)]=  matrix_mx
    
    
        intervened_times = np.zeros(graph.num_parents_y)
        intervened_times_list[str(m)]=   intervened_times
        by = np.array([np.zeros(graph.num_parents_y)]).T
        by_list[str(m)]= by
        payoff_list = [0]
        total_expected_payoff = 0
        total_optimal_reward=0
    
    # we have a graph list
    #

    T_remaining=graph.T-T_initial
    for i in range(T_initial):
        for m in range(len(graph_list)):
            graph=graph_list[m]
            best_intervention=[0,1,2]
            total_expected_payoff += graph.expect_y(best_intervention)
            # total_expected_payoff+=0
            total_optimal_reward+=best_y_list[m]

        

    
    for i in range(T_remaining):
        for m in range(len(graph_list)):
            graph=graph_list[m]
            # matrix_m=  matrix_m_list[str(m)]

            matrix_mx=  matrix_mx_list[str(m)]
            matrix_m_one=matrix_m_one_list[str(m)]
            matrix_m_two=matrix_m_two_list[str(m)]
            by_one= by_one_list[str(m)]
            by_two= by_two_list[str(m)]

            intervened_times=intervened_times_list[str(m)]
            by= by_list[str(m)] 

            # inverse_m = np.linalg.inv(matrix_m)
            # print("matrix_m:",  matrix_m.shape)
            # inverse_m_one=np.linalg.inv(matrix_m_one)
            
            inverse_m_one = np.linalg.inv(matrix_m_one)
            inverse_m_two = np.linalg.inv(matrix_m_two)



            hat_theta_x = np.zeros(graph.num_parents_y)
            for j in range(graph.num_parents_y):
                hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
            hat_theta_x = np.array([hat_theta_x]).T
            norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])

        
            
            hat_theta_y_one = np.matmul(inverse_m_one, np.array([by_one]).T)
            hat_theta_y_two = np.matmul(inverse_m_two, np.array([by_two]).T)
            
        
            rho = graph.compute_rho_lr(i + 1)

            best_intervention = find_best_intervention_cluster_y(graph,inverse_m_one,inverse_m_two, hat_theta_x,hat_theta_y_one,hat_theta_y_two,norm_matrix_mx, rho)
            
            vy, y= graph.simulate(best_intervention)
        
            vy_g_one=vy[:3]
            vy_g_two=vy[-3:]

            y_g_one=y-np.dot(hat_theta_y_two.flatten(),vy_g_two.flatten())
            y_g_two=y-np.dot(hat_theta_y_one.flatten(),vy_g_one.flatten())
            

            
            total_expected_payoff += graph.expect_y(best_intervention)
            total_optimal_reward+=best_y_list[m]


            if (m%2==0):
                matrix_m_one_sum+= np.matmul(vy_g_one,vy_g_one.T)
                by_one_sum+=y_g_one*vy_g_one
                matrix_m_two_sum+= np.matmul(vy_g_two,vy_g_two.T)
                by_two_sum+=y_g_two*vy_g_two
                # print("m_one",m)
            
            else:
                matrix_m_one_log+= np.matmul(vy_g_one,vy_g_one.T)
                by_one_log+=y_g_one*vy_g_one
                matrix_m_two_log+= np.matmul(vy_g_two,vy_g_two.T)
                by_two_log+=y_g_two*vy_g_two
                
                # print("m_two":m)
            
            
            # matrix_m += np.matmul(vy, vy.T)
            # by += y * vy
            for j in best_intervention:
                intervened_times[int(j)] += 1
            for j in range(graph.num_parents_y):
                if j not in best_intervention:
                    matrix_mx[j] += vy[j]

            matrix_m_one_list[str(m)]=  matrix_m_one 
            matrix_m_two_list[str(m)]=  matrix_m_two
            by_one_list[str(m)]= by_one
            by_two_list[str(m)]= by_two

            matrix_mx_list[str(m)]=matrix_mx
            intervened_times_list[str(m)]=intervened_times
        
        for m in range(len(graph_list)):
            if (m%2==0):
                # print("m_one_agg:",m)
                matrix_m_one_list[str(m)]= matrix_m_one_sum
                by_one_list[str(m)]=by_one_sum
                matrix_m_two_list[str(m)]= matrix_m_two_sum
                by_two_list[str(m)]=by_two_sum

            else:
                # print("m_one_agg:",m)
                matrix_m_one_list[str(m)]= matrix_m_one_log
                by_one_list[str(m)]=by_one_log
                matrix_m_two_list[str(m)]= matrix_m_two_log
                by_two_list[str(m)]=by_two_log
            
            


        if i % 100 == 99:
            print("partition_vec")
            
            print("round:",i)

            print("reward:", total_expected_payoff)
            total_regret=total_optimal_reward-total_expected_payoff
            payoff_list.append(total_regret)

            
        
    return payoff_list, total_expected_payoff
# def homogeneity_testing(design_matrix,theta_list):


def independent_explore(graph_list):
    NoiseScale=100
    graph=graph_list[0]
    T_0=1000
    
    vy_two_list=[]

    client_design_list={}
    client_theta_list={}
    client_cov_list={}
    client_by_list={}
  

    intervention_now=[0,1,2]
    
    # intervention_now=[3,4,5]
    for m in range(len(graph_list)):
        graph=graph_list[m]
        vy_one_list=[]
        matrix_m_one = np.zeros((3, 3))
        by_one = np.array([np.zeros(3)]).T
        vy_one_arr=np.zeros((T_0,3,1))
        for i in range(3):
            matrix_m_one[i][i] = 1
        for i in range(T_0):
            vy,y_one=graph.simulate_explore(intervention_now)
            if (intervention_now==[0,1,2]):
                vy_one=vy[-3:]
            if (intervention_now==[3,4,5]):
                vy_one=vy[:3]
            vy_one_list.append(vy_one)
            vy_one_arr=np.array(vy_one_list)
            matrix_m_one+= np.matmul( vy_one, vy_one.T)
            by_one+=y_one* vy_one
        inverse_m_one = np.linalg.inv(matrix_m_one)
        hat_theta_y_one = np.matmul(inverse_m_one, np.array([by_one]).T)
        vy_one_arr=vy_one_arr.squeeze(-1)
        client_cov_list[str(m)]= matrix_m_one
        client_by_list[str(m)]=  by_one
        client_design_list[str(m)]= vy_one_arr
        client_theta_list[str(m)]=hat_theta_y_one

    for m in range(len(graph_list)):
        print("current:",m)
        hat_theta_y_one=client_theta_list[str(m)]
        vy_one_arr=client_design_list[str(m)]
        matrix_m_one= client_cov_list[str(m)]
        by_one=  client_by_list[str(m)]
        for n in range(len(graph_list)):
            if (n!=m):
                print("neighbor:",n)
                matrix_m_two= client_cov_list[str(n)]
                by_two=  client_by_list[str(n)]
                hat_theta_y_two=client_theta_list[str(n)]
                vy_two_arr=client_design_list[str(n)]

                inverse_m_combine=np.linalg.inv(matrix_m_one+matrix_m_two)
                by_combine=by_one+by_two
                hat_theta_y_combine = np.matmul(inverse_m_combine, np.array([by_combine]).T)

                theta_difference_one=hat_theta_y_one - hat_theta_y_combine
                theta_difference_two=hat_theta_y_two - hat_theta_y_combine

                theta_difference_one=  theta_difference_one.squeeze(-1).T
                theta_difference_two=  theta_difference_two.squeeze(-1).T

                num = np.linalg.norm(np.matmul(vy_one_arr,    theta_difference_one )) ** 2 + np.linalg.norm(
                    np.matmul( vy_two_arr,  theta_difference_two)) ** 2

                XCombinedRank = np.linalg.matrix_rank(np.concatenate((vy_one_arr,vy_two_arr), axis=0))
                df1 = int(vy_one_arr.ndim +vy_two_arr.ndim - XCombinedRank)
                chiSquareStatistic = num / (NoiseScale**2)
                psi = 1 / (NoiseScale**2* len(graph_list)**2)
                p_value = ncx2.sf(x=chiSquareStatistic, df=df1, nc=psi)
                print("  p_value:",  p_value)
        print("______________________")





   

   

    # for i in range(T_0):
    #     intervention_now=[0,1,2]
    #     vy,y_two=graph.simulate_explore(intervention_now)
    #     vy_two=vy[-3:]
    #     vy_two_list.append(vy_two)
    #     vy_two_arr=np.array(vy_two_list)
    #     matrix_m_two+= np.matmul( vy_two, vy_two.T)
    #     by_two+=y_two* vy_two

    # inverse_m_two = np.linalg.inv(matrix_m_two)
    # hat_theta_y_two = np.matmul(inverse_m_two, np.array([by_two]).T)
    # vy_two_arr=vy_two_arr.squeeze(-1)

    # inverse_m_combine=np.linalg.inv(matrix_m_one+matrix_m_two)
    # by_combine=by_one+by_two
    # hat_theta_y_combine = np.matmul(inverse_m_combine, np.array([by_combine]).T)

   

def run_blm_lr_partition_cluster(graph_list,best_y_list):
    graph=graph_list[0]
    matrix_m_one_list={}
    matrix_m_two_list={}
    by_one_list={}
    by_two_list={}
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}

    
    for m in range(len(graph_list)):
        matrix_m_one = np.zeros((3, 3))
        matrix_m_two = np.zeros((3, 3))
        
        
        for i in range(3):
            matrix_m_one[i][i] =lambda_matrix
            matrix_m_two[i][i] = lambda_matrix
        by_one = np.array([np.zeros(3)]).T
        by_two = np.array([np.zeros(3)]).T

        matrix_m_one_sum= np.zeros((3, 3))
        matrix_m_two_sum= np.zeros((3, 3))
        for i in range(3):
            matrix_m_one_sum[i][i] = lambda_matrix
            matrix_m_two_sum[i][i] = lambda_matrix
        by_one_sum = np.array([np.zeros(3)]).T
        by_two_sum = np.array([np.zeros(3)]).T

        matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
        for i in range(graph.num_parents_y):
            matrix_m[i][i] = lambda_matrix

        
    
        matrix_m_one_list[str(m)]=  matrix_m_one 
        matrix_m_two_list[str(m)]=  matrix_m_two
        by_one_list[str(m)]= by_one
        by_two_list[str(m)]= by_two

        matrix_m_list[str(m)]=  matrix_m
        matrix_mx = np.zeros(graph.num_parents_y)
        matrix_mx_list[str(m)]=  matrix_mx
    
    
        intervened_times = np.zeros(graph.num_parents_y)
        intervened_times_list[str(m)]=   intervened_times
        by = np.array([np.zeros(graph.num_parents_y)]).T
        by_list[str(m)]= by
        payoff_list = [0]
        total_expected_payoff = 0
        total_optimal_reward=0
    
    # we have a graph list
    #
    for i in range(graph.T):
        
        for m in range(len(graph_list)):
            graph=graph_list[m]
            # matrix_m=  matrix_m_list[str(m)]

            matrix_mx=  matrix_mx_list[str(m)]
            matrix_m_one=matrix_m_one_list[str(m)]
            matrix_m_two=matrix_m_two_list[str(m)]
            by_one= by_one_list[str(m)]
            by_two= by_two_list[str(m)]

            intervened_times=intervened_times_list[str(m)]
            by= by_list[str(m)] 

            # inverse_m = np.linalg.inv(matrix_m)
            # print("matrix_m:",  matrix_m.shape)
            # inverse_m_one=np.linalg.inv(matrix_m_one)
            
            inverse_m_one = np.linalg.inv(matrix_m_one)
            inverse_m_two = np.linalg.inv(matrix_m_two)



            hat_theta_x = np.zeros(graph.num_parents_y)
            for j in range(graph.num_parents_y):
                hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
            hat_theta_x = np.array([hat_theta_x]).T
            norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])

        
            
            hat_theta_y_one = np.matmul(inverse_m_one, np.array([by_one]).T)
        
            hat_theta_y_two = np.matmul(inverse_m_two, np.array([by_two]).T)
            
            # hat_theta_y_con = np.concatenate((hat_theta_y_one,hat_theta_y_two), axis=1)
            
            # hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
            # print("hat_theta_y:",hat_theta_y)
            rho = graph.compute_rho_lr(i + 1)

            best_intervention = find_best_intervention_cluster_y(graph,inverse_m_one,inverse_m_two, hat_theta_x,hat_theta_y_one,hat_theta_y_two,norm_matrix_mx, rho)
            
            vy, y= graph.simulate(best_intervention)
        
            vy_g_one=vy[:3]
            vy_g_two=vy[-3:]
            
            y_g_one=y-np.dot(hat_theta_y_two.flatten(),vy_g_two.flatten())
            y_g_two=y-np.dot(hat_theta_y_one.flatten(),vy_g_one.flatten())
            
            
            total_expected_payoff += graph.expect_y(best_intervention)
            total_optimal_reward+=best_y_list[m]

            if (m%2==0):
                matrix_m_one_sum+= np.matmul(vy_g_one,vy_g_one.T)
                by_one_sum+=y_g_one*vy_g_one
                matrix_m_two+= np.matmul(vy_g_two,vy_g_two.T)
                by_two+=y_g_two*vy_g_two
                # print("m_one",m)
            
            else:
                matrix_m_one+= np.matmul(vy_g_one,vy_g_one.T)
                by_one+=y_g_one*vy_g_one
                matrix_m_two_sum+= np.matmul(vy_g_two,vy_g_two.T)
                by_two_sum+=y_g_two*vy_g_two
                # print("m_two":m)
            

           
            
            
            # matrix_m += np.matmul(vy, vy.T)
            # by += y * vy
            for j in best_intervention:
                intervened_times[int(j)] += 1
            for j in range(graph.num_parents_y):
                if j not in best_intervention:
                    matrix_mx[j] += vy[j]

            matrix_m_one_list[str(m)]=  matrix_m_one 
            matrix_m_two_list[str(m)]=  matrix_m_two
            by_one_list[str(m)]= by_one
            by_two_list[str(m)]= by_two

            matrix_mx_list[str(m)]=matrix_mx
            intervened_times_list[str(m)]=intervened_times
        
        for m in range(len(graph_list)):
            if (m%2==0):
                # print("m_one_agg:",m)
                matrix_m_one_list[str(m)]= matrix_m_one_sum
                by_one_list[str(m)]=by_one_sum
            else:
                # print("m_two_agg:",m)
                matrix_m_two_list[str(m)]= matrix_m_two_sum
                by_two_list[str(m)]=by_two_sum


        if i % 100 == 99:
            print("partition_c")
            
            print("round:",i)
            total_regret=  total_optimal_reward-total_expected_payoff

            print("reward:", total_expected_payoff)
            print("regret:",total_regret)
            payoff_list.append(total_regret)

            
    return payoff_list, total_expected_payoff

def run_blm_lr_partition_cluster_two(graph_list,best_y_list):
    graph=graph_list[0]
    matrix_m_one_list={}
    matrix_m_two_list={}
    matrix_m_three_list={}
    by_one_list={}
    by_two_list={}
    by_three_list={}
    matrix_mx_list={}
    intervened_times_list={}


    
    for m in range(len(graph_list)):
        matrix_m_one = np.zeros((2, 2))
        matrix_m_two = np.zeros((2, 2))
        
        
        for i in range(2):
            matrix_m_one[i][i] =lambda_matrix
            matrix_m_two[i][i] = lambda_matrix
        by_one = np.array([np.zeros(2)]).T
        by_two = np.array([np.zeros(2)]).T

        matrix_m_one_sum= np.zeros((2, 2))
        matrix_m_two_sum= np.zeros((2, 2))
        for i in range(2):
            matrix_m_one_sum[i][i] = lambda_matrix
            matrix_m_two_sum[i][i] = lambda_matrix
        by_one_sum = np.array([np.zeros(2)]).T
        by_two_sum = np.array([np.zeros(2)]).T

   

        
    
        matrix_m_one_list[str(m)]=  matrix_m_one 
        matrix_m_two_list[str(m)]=  matrix_m_two
        matrix_m_three_list[str(m)]=matrix_m_one 
        by_one_list[str(m)]= by_one
        by_two_list[str(m)]= by_two
        by_three_list[str(m)]=by_one

       
        matrix_mx = np.zeros(graph.num_parents_y)
        matrix_mx_list[str(m)]=  matrix_mx
    
    
        intervened_times = np.zeros(graph.num_parents_y)
        intervened_times_list[str(m)]=   intervened_times
        payoff_list = [0]
        total_expected_payoff = 0
        total_optimal_reward=0
    
    # we have a graph list
    #
    T_remaining=graph.T-T_initial
    for i in range(T_initial):
        for m in range(len(graph_list)):
            graph=graph_list[m]
            best_intervention=[0,1,2]
            total_expected_payoff += graph.expect_y(best_intervention)
            # total_expected_payoff+=0
            total_optimal_reward+=best_y_list[m]

        

    
    for i in range(T_remaining):
        
        for m in range(len(graph_list)):
            graph=graph_list[m]
            # matrix_m=  matrix_m_list[str(m)]

            matrix_mx=  matrix_mx_list[str(m)]
            matrix_m_one=matrix_m_one_list[str(m)]
            matrix_m_two=matrix_m_two_list[str(m)]
            matrix_m_three=matrix_m_three_list[str(m)]
            by_one= by_one_list[str(m)]
            by_two= by_two_list[str(m)]
            by_three= by_three_list[str(m)]

            intervened_times=intervened_times_list[str(m)]
         
            
            inverse_m_one = np.linalg.inv(matrix_m_one)
            inverse_m_two = np.linalg.inv(matrix_m_two)
            inverse_m_three = np.linalg.inv( matrix_m_three)



            hat_theta_x = np.zeros(graph.num_parents_y)
            for j in range(graph.num_parents_y):
                hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
            hat_theta_x = np.array([hat_theta_x]).T
            norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])

        
            
            hat_theta_y_one = np.matmul(inverse_m_one, np.array([by_one]).T)
        
            hat_theta_y_two = np.matmul(inverse_m_two, np.array([by_two]).T)
            
            hat_theta_y_three = np.matmul(inverse_m_three, np.array([by_three]).T)
            
            # hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
            # print("hat_theta_y:",hat_theta_y)
            rho = graph.compute_rho_lr(i + 1)

            best_intervention = find_best_intervention_cluster_two(graph,inverse_m_one,inverse_m_two, inverse_m_three, hat_theta_x,hat_theta_y_one,hat_theta_y_two,hat_theta_y_three,norm_matrix_mx, rho)
            
            vy, y_list= graph.simulate_two(best_intervention)
        
            vy_g_one=vy[:2]
            vy_g_two=vy[-2:]
            vy_g_three=vy[2:4]
        
            
            
            total_expected_payoff += graph.expect_y(best_intervention)
            total_optimal_reward+=best_y_list[m]

            matrix_m_three+= np.matmul(vy_g_three,vy_g_three.T)
            by_three+=y_list[2]*vy_g_three

            if (m%2==0):
                matrix_m_one_sum+= np.matmul(vy_g_one,vy_g_one.T)
                by_one_sum+=y_list[0]*vy_g_one
                matrix_m_two+= np.matmul(vy_g_two,vy_g_two.T)
                by_two+=y_list[1]*vy_g_two
               
                # print("m_one",m)
            
            else:
                matrix_m_one+= np.matmul(vy_g_one,vy_g_one.T)
                by_one+=y_list[0]*vy_g_one
                matrix_m_two_sum+= np.matmul(vy_g_two,vy_g_two.T)
                by_two_sum+=y_list[1]*vy_g_two
                # print("m_two":m)
            

           
            
            
            # matrix_m += np.matmul(vy, vy.T)
            # by += y * vy
            for j in best_intervention:
                intervened_times[int(j)] += 1
            for j in range(graph.num_parents_y):
                if j not in best_intervention:
                    matrix_mx[j] += vy[j]

            matrix_m_one_list[str(m)]=  matrix_m_one 
            matrix_m_two_list[str(m)]=  matrix_m_two
            matrix_m_three_list[str(m)]=  matrix_m_three
            by_one_list[str(m)]= by_one
            by_two_list[str(m)]= by_two
            by_three_list[str(m)]= by_three

            matrix_mx_list[str(m)]=matrix_mx
            intervened_times_list[str(m)]=intervened_times
        
        for m in range(len(graph_list)):
            if (m%2==0):
                # print("m_one_agg:",m)
                matrix_m_one_list[str(m)]= matrix_m_one_sum
                by_one_list[str(m)]=by_one_sum
            else:
                # print("m_two_agg:",m)
                matrix_m_two_list[str(m)]= matrix_m_two_sum
                by_two_list[str(m)]=by_two_sum


        if i % 100 == 99:
            print("partition_ctwo")
            
            print("round:",i)
            total_regret=  total_optimal_reward-total_expected_payoff

            print("reward:", total_expected_payoff)
            print("regret:",total_regret)
            payoff_list.append(total_regret)

            
    return payoff_list, total_expected_payoff


def run_blm_lr_partition_cluster_component(graph_list,best_y_list):
    graph=graph_list[0]
    # 创建一个空字典，用于存储多个列表
    matrix_m_lists_dict = {}
    by_lists_dict = {}

    # 循环创建多个列表，并以 str(i) 作为键
    for i in range(6): 
        matrix_m_lists_dict[f"matrix_m_list_{i}"] = {}  
        by_lists_dict[f"by_m_list_{i}"] = {}   


    
    matrix_mx_list={}
    intervened_times_list={}
 

    
    for m in range(len(graph_list)):
        matrix_m_one = np.zeros((1, 1))
        for i in range(1):
            matrix_m_one[i][i] =lambda_matrix
           
        by_one = np.array([np.zeros(1)]).T

        for j in range(6):
            matrix_m_lists_dict[f"matrix_m_list_{j}"][str(m)]=matrix_m_one
            by_lists_dict[f"by_m_list_{j}"][str(m)]=by_one

   
      
        matrix_sum_dict = {}
        for i in range(6):  
            matrix_sum_dict[f"matrix_sum_{i}"] =np.zeros((1, 1)) 
            for j in range(1):
                matrix_sum_dict[f"matrix_sum_{i}"][j][j] = lambda_matrix
        
        by_sum_dict = {}
        for i in range(6):  
            by_sum_dict[f"by_sum_{i}"] =np.array([np.zeros(1)]).T 
       

        matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
        for i in range(graph.num_parents_y):
            matrix_m[i][i] = lambda_matrix

        

        matrix_mx = np.zeros(graph.num_parents_y)
        matrix_mx_list[str(m)]=  matrix_mx
    
    
        intervened_times = np.zeros(graph.num_parents_y)
        intervened_times_list[str(m)]=   intervened_times
        
        payoff_list = [0]
        total_expected_payoff = 0
        total_optimal_reward=0
    
    # we have a graph list
    #
    T_remaining=graph.T-T_initial
    for i in range(T_initial):
        for m in range(len(graph_list)):
            graph=graph_list[m]
            best_intervention=[0,1,2]
            total_expected_payoff += graph.expect_y(best_intervention)
            # total_expected_payoff+=0
            total_optimal_reward+=best_y_list[m]

        

    
    for i in range(T_remaining):
        
        for m in range(len(graph_list)):
            graph=graph_list[m]
            # matrix_m=  matrix_m_list[str(m)]

            matrix_mx=  matrix_mx_list[str(m)]
            intervened_times=intervened_times_list[str(m)]




            hat_theta_x = np.zeros(graph.num_parents_y)
            for j in range(graph.num_parents_y):
                hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
            hat_theta_x = np.array([hat_theta_x]).T
            norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])

            hat_theta_y_list=[]
            inverse_m_list=[]

            for j in range(6):
                matrix_m_now=matrix_m_lists_dict[f"matrix_m_list_{j}"][str(m)]
                inverse_m_list.append(matrix_m_now)
                by_now=by_lists_dict[f"by_m_list_{j}"][str(m)]

                hat_theta_y=np.matmul(matrix_m_now, np.array([by_now]).T)
                hat_theta_y_list.append(  hat_theta_y)
            
            rho = graph.compute_rho_lr(i + 1)

            best_intervention = find_best_intervention_cluster_component(graph,inverse_m_list, hat_theta_x,hat_theta_y_list,norm_matrix_mx, rho)
            
            vy, y_list= graph.simulate_component(best_intervention)
            # print("y_list",y_list)

            for j in range(6):
                vy_g_now=vy[j]
                y_g_now=y_list[j]
                matrix_sum_dict[f"matrix_sum_{j}"]+=np.matmul(vy_g_now,vy_g_now.T)
                by_sum_dict[f"by_sum_{j}"] += y_g_now*  vy_g_now
        
            
            
            
            total_expected_payoff += graph.expect_y(best_intervention)
            total_optimal_reward+=best_y_list[m]

         

           
            
            
            # matrix_m += np.matmul(vy, vy.T)
            # by += y * vy
            for j in best_intervention:
                intervened_times[int(j)] += 1
            for j in range(graph.num_parents_y):
                if j not in best_intervention:
                    matrix_mx[j] += vy[j]

            matrix_mx_list[str(m)]=matrix_mx
            intervened_times_list[str(m)]=intervened_times
        
        for m in range(len(graph_list)):
            for j in range(6):
                matrix_m_lists_dict[f"matrix_m_list_{j}"][str(m)]=matrix_sum_dict[f"matrix_sum_{j}"]
                by_lists_dict[f"by_m_list_{j}"][str(m)]=by_sum_dict[f"by_sum_{j}"]




        if i % 100 == 99:
            print("partition_component")
            
            print("round:",i)
            total_regret=  total_optimal_reward-total_expected_payoff

            print("reward:", total_expected_payoff)
            print("regret:",total_regret)
            payoff_list.append(total_regret)

            
    return payoff_list, total_expected_payoff

def run_blm_lr_partition_cluster_explore(graph_list,best_y_list):
    graph=graph_list[0]
    matrix_m_one_list={}
    matrix_m_two_list={}
    by_one_list={}
    by_two_list={}
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}

    
    for m in range(len(graph_list)):
        matrix_m_one = np.zeros((3, 3))
        matrix_m_two = np.zeros((3, 3))
        
        
        for i in range(3):
            matrix_m_one[i][i] =lambda_matrix
            matrix_m_two[i][i] = lambda_matrix
        by_one = np.array([np.zeros(3)]).T
        by_two = np.array([np.zeros(3)]).T

        matrix_m_one_sum= np.zeros((3, 3))
        matrix_m_two_sum= np.zeros((3, 3))
        for i in range(3):
            matrix_m_one_sum[i][i] = lambda_matrix
            matrix_m_two_sum[i][i] = lambda_matrix
        by_one_sum = np.array([np.zeros(3)]).T
        by_two_sum = np.array([np.zeros(3)]).T

        matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
        for i in range(graph.num_parents_y):
            matrix_m[i][i] = lambda_matrix

        
    
        matrix_m_one_list[str(m)]=  matrix_m_one 
        matrix_m_two_list[str(m)]=  matrix_m_two
        by_one_list[str(m)]= by_one
        by_two_list[str(m)]= by_two

        matrix_m_list[str(m)]=  matrix_m
        matrix_mx = np.zeros(graph.num_parents_y)
        matrix_mx_list[str(m)]=  matrix_mx
    
    
        intervened_times = np.zeros(graph.num_parents_y)
        intervened_times_list[str(m)]=   intervened_times
        by = np.array([np.zeros(graph.num_parents_y)]).T
        by_list[str(m)]= by
        payoff_list = [0]
        total_expected_payoff = 0
        total_optimal_reward=0
    
    # we have a graph list
    #
 
    T_remaining=graph.T-T_initial
    for i in range(T_initial):
        for m in range(len(graph_list)):
            graph=graph_list[m]
            best_intervention=[0,1,2]
            total_expected_payoff += graph.expect_y(best_intervention)
            # total_expected_payoff+=0
            total_optimal_reward+=best_y_list[m]

       

    
    for i in range(T_remaining):
        
        for m in range(len(graph_list)):
            graph=graph_list[m]
            # matrix_m=  matrix_m_list[str(m)]

            matrix_mx=  matrix_mx_list[str(m)]
            matrix_m_one=matrix_m_one_list[str(m)]
            matrix_m_two=matrix_m_two_list[str(m)]
            by_one= by_one_list[str(m)]
            by_two= by_two_list[str(m)]

            intervened_times=intervened_times_list[str(m)]
            by= by_list[str(m)] 

            # inverse_m = np.linalg.inv(matrix_m)
            # print("matrix_m:",  matrix_m.shape)
            # inverse_m_one=np.linalg.inv(matrix_m_one)
            
            inverse_m_one = np.linalg.inv(matrix_m_one)
            inverse_m_two = np.linalg.inv(matrix_m_two)



            hat_theta_x = np.zeros(graph.num_parents_y)
            for j in range(graph.num_parents_y):
                hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
            hat_theta_x = np.array([hat_theta_x]).T
            norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])

        
            
            hat_theta_y_one = np.matmul(inverse_m_one, np.array([by_one]).T)
        
            hat_theta_y_two = np.matmul(inverse_m_two, np.array([by_two]).T)
            
            # hat_theta_y_con = np.concatenate((hat_theta_y_one,hat_theta_y_two), axis=1)
            
            # hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
            # print("hat_theta_y:",hat_theta_y)
            rho = graph.compute_rho_lr(i + 1)

            best_intervention = find_best_intervention_cluster_y(graph,inverse_m_one,inverse_m_two, hat_theta_x,hat_theta_y_one,hat_theta_y_two,norm_matrix_mx, rho)
            
            vy, y= graph.simulate(best_intervention)
        
            vy_g_one=vy[:3]
            vy_g_two=vy[-3:]
            
            y_g_one=y-np.dot(hat_theta_y_two.flatten(),vy_g_two.flatten())
            y_g_two=y-np.dot(hat_theta_y_one.flatten(),vy_g_one.flatten())
            
            
        
            total_expected_payoff += graph.expect_y(best_intervention)
            total_optimal_reward+=best_y_list[m]

            if (m%2==0):
                matrix_m_one_sum+= np.matmul(vy_g_one,vy_g_one.T)
                by_one_sum+=y_g_one*vy_g_one
                matrix_m_two+= np.matmul(vy_g_two,vy_g_two.T)
                by_two+=y_g_two*vy_g_two
                # print("m_one",m)
            
            else:
                matrix_m_one+= np.matmul(vy_g_one,vy_g_one.T)
                by_one+=y_g_one*vy_g_one
                matrix_m_two_sum+= np.matmul(vy_g_two,vy_g_two.T)
                by_two_sum+=y_g_two*vy_g_two
             
            

            for j in best_intervention:
                intervened_times[int(j)] += 1
            for j in range(graph.num_parents_y):
                if j not in best_intervention:
                    matrix_mx[j] += vy[j]

            matrix_m_one_list[str(m)]=  matrix_m_one 
            matrix_m_two_list[str(m)]=  matrix_m_two
            by_one_list[str(m)]= by_one
            by_two_list[str(m)]= by_two

            matrix_mx_list[str(m)]=matrix_mx
            intervened_times_list[str(m)]=intervened_times
        
        for m in range(len(graph_list)):
            if (m%2==0):
                # print("m_one_agg:",m)
                matrix_m_one_list[str(m)]= matrix_m_one_sum
                by_one_list[str(m)]=by_one_sum
            else:
                # print("m_two_agg:",m)
                matrix_m_two_list[str(m)]= matrix_m_two_sum
                by_two_list[str(m)]=by_two_sum


        if i % 100 == 99:
            print("partition_c")
            
            print("round:",i)
            total_regret=  total_optimal_reward-total_expected_payoff

            print("reward:", total_expected_payoff)
            print("regret:",total_regret)
            payoff_list.append(total_regret)

            
        
    return payoff_list, total_expected_payoff

def run_blm_lr_partition_Ind(graph_list,best_y_list):
    graph=graph_list[0]
    matrix_m_one_list={}
    matrix_m_two_list={}
    by_one_list={}
    by_two_list={}
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}

    if graph.type == 'parallel':
        for m in range(len(graph_list)):
            matrix_m_one = np.zeros((3, 3))
            matrix_m_two = np.zeros((3, 3))
            
          
            for i in range(3):
                matrix_m_one[i][i] = lambda_matrix
                matrix_m_two[i][i] =lambda_matrix
            by_one = np.array([np.zeros(3)]).T
            by_two = np.array([np.zeros(3)]).T

            matrix_m_one_sum= np.zeros((3, 3))
            matrix_m_two_sum= np.zeros((3, 3))
            by_one_sum = np.array([np.zeros(3)]).T
            by_two_sum = np.array([np.zeros(3)]).T

            matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
            for i in range(graph.num_parents_y):
                matrix_m[i][i] =lambda_matrix

            
        
            matrix_m_one_list[str(m)]=  matrix_m_one 
            matrix_m_two_list[str(m)]=  matrix_m_two
            by_one_list[str(m)]= by_one
            by_two_list[str(m)]= by_two

            matrix_m_list[str(m)]=  matrix_m
            matrix_mx = np.zeros(graph.num_parents_y)
            matrix_mx_list[str(m)]=  matrix_mx
      
        
            intervened_times = np.zeros(graph.num_parents_y)
            intervened_times_list[str(m)]=   intervened_times
            by = np.array([np.zeros(graph.num_parents_y)]).T
            by_list[str(m)]= by
            payoff_list = [0]
            total_expected_payoff = 0
        
        # we have a graph list
        #
        for i in range(graph.T):
            for m in range(len(graph_list)):
                graph=graph_list[m]
                # matrix_m=  matrix_m_list[str(m)]

                matrix_mx=  matrix_mx_list[str(m)]
                matrix_m_one=matrix_m_one_list[str(m)]
                matrix_m_two=matrix_m_two_list[str(m)]
                by_one= by_one_list[str(m)]
                by_two= by_two_list[str(m)]

                intervened_times=intervened_times_list[str(m)]
                by= by_list[str(m)] 

                # inverse_m = np.linalg.inv(matrix_m)
                # print("matrix_m:",  matrix_m.shape)
                # inverse_m_one=np.linalg.inv(matrix_m_one)
               
                inverse_m_one = np.linalg.inv(matrix_m_one)
                inverse_m_two = np.linalg.inv(matrix_m_two)



                hat_theta_x = np.zeros(graph.num_parents_y)
                for j in range(graph.num_parents_y):
                    hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
                hat_theta_x = np.array([hat_theta_x]).T
                norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])

            
             
                hat_theta_y_one = np.matmul(inverse_m_one, np.array([by_one]).T)
                hat_theta_y_two = np.matmul(inverse_m_two, np.array([by_two]).T)
              
                # hat_theta_y_con = np.concatenate((hat_theta_y_one,hat_theta_y_two), axis=1)
                
                # hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
                # print("hat_theta_y:",hat_theta_y)
                rho = graph.compute_rho_lr(i + 1)

                best_intervention = find_best_intervention_cluster_y(graph,inverse_m_one,inverse_m_two, hat_theta_x,hat_theta_y_one,hat_theta_y_two,norm_matrix_mx, rho)
                
                vy, y,y_g_one,y_g_two= graph.simulate_cluster_y(best_intervention)
          
                vy_g_one=vy[:3]
                vy_g_two=vy[-3:]
               

                total_expected_payoff += graph.expect_y(best_intervention)
                total_optimal_reward+=best_y_list[m]

                matrix_m_one+= np.matmul(vy_g_one,vy_g_one.T)
                matrix_m_two+= np.matmul(vy_g_two,vy_g_two.T)
                by_one+=y_g_one*vy_g_one
                by_two+=y_g_two*vy_g_two
               
                # matrix_m += np.matmul(vy, vy.T)
                # by += y * vy
                for j in best_intervention:
                    intervened_times[int(j)] += 1
                for j in range(graph.num_parents_y):
                    if j not in best_intervention:
                        matrix_mx[j] += vy[j]

                matrix_m_one_list[str(m)]=  matrix_m_one 
                matrix_m_two_list[str(m)]=  matrix_m_two
                by_one_list[str(m)]= by_one
                by_two_list[str(m)]= by_two

                matrix_mx_list[str(m)]=matrix_mx
                intervened_times_list[str(m)]=intervened_times

            if i % 100 == 99:
                print("partition_Ind")
              
                print("round:",i)

                print("reward:", total_expected_payoff)
                total_regret= total_optimal_reward-total_expected_payoff
                payoff_list.append(total_regret)

                
            # print(best_intervention)
        # print(np.matmul(np.linalg.inv(matrix_m), np.array([by]).T))
        return payoff_list, total_expected_payoff
   
def run_blm_lr_cluster_x(graph_list):
    graph=graph_list[0]
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}

    if graph.type == 'parallel':
        for m in range(len(graph_list)):
            matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
            for i in range(graph.num_parents_y):
                matrix_m[i][i] = 1
           
            matrix_m_list[str(m)]=  matrix_m
            matrix_mx = np.zeros(graph.num_parents_y)
            matrix_mx_list[str(m)]=  matrix_mx
      
        
            intervened_times = np.zeros(graph.num_parents_y)
            intervened_times_list[str(m)]=   intervened_times
            by = np.array([np.zeros(graph.num_parents_y)]).T
            by_list[str(m)]= by
            payoff_list = [0]
            total_expected_payoff = 0
        
        # we have a graph list
        #
        for i in range(graph.T):
            for m in range(len(graph_list)):
                graph=graph_list[m]
                matrix_m=  matrix_m_list[str(m)]
                matrix_mx=  matrix_mx_list[str(m)]

                intervened_times=intervened_times_list[str(m)]
                by= by_list[str(m)] 



                inverse_m = np.linalg.inv(matrix_m)
                hat_theta_x = np.zeros(graph.num_parents_y)
                for j in range(graph.num_parents_y):
                    hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
                hat_theta_x = np.array([hat_theta_x]).T
                norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])
                hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
                rho = graph.compute_rho_lr(i + 1)
                best_intervention = find_best_intervention(graph, inverse_m, hat_theta_x, hat_theta_y, norm_matrix_mx, rho)
                vy, y = graph.simulate(best_intervention)
                total_expected_payoff += graph.expect_y(best_intervention)
                
               
               
                matrix_m += np.matmul(vy, vy.T)
                by += y * vy
                for j in best_intervention:
                    intervened_times[int(j)] += 1
                for j in range(graph.num_parents_y):
                    if j not in best_intervention:
                        matrix_mx[j] += vy[j]

                matrix_m_list[str(m)]= matrix_m
                matrix_mx_list[str(m)]=matrix_mx
                intervened_times_list[str(m)]=intervened_times
                by_list[str(m)] =by

            for m in range(len(graph_list)):
                if (m==0):
                    for j in range(graph.num_parents_y):
                        if(j<3):
                            matrix_mx_list[str(m)][j]=  matrix_mx_list[str(m)][j]+ matrix_mx_list[str(1)][j]
                            intervened_times_list[str(m)][j]= intervened_times_list[str(m)][j]+ intervened_times_list[str(1)][j]

                        if(j>2):
                            matrix_mx_list[str(m)][j]=  matrix_mx_list[str(m)][j]+ matrix_mx_list[str(2)][j]
                            intervened_times_list[str(m)][j]= intervened_times_list[str(m)][j]+ intervened_times_list[str(2)][j]
                if (m==1):
                    for j in range(graph.num_parents_y):
                        if(j<3):
                            matrix_mx_list[str(m)][j]=  matrix_mx_list[str(0)][j]
                            intervened_times_list[str(m)][j]= intervened_times_list[str(0)][j]

                        if(j>2):
                            matrix_mx_list[str(m)][j]=  matrix_mx_list[str(m)][j]+ matrix_mx_list[str(3)][j]
                            intervened_times_list[str(m)][j]= intervened_times_list[str(m)][j]+ intervened_times_list[str(3)][j]

                if (m==2):
                    for j in range(graph.num_parents_y):
                        if(j<3):
                            matrix_mx_list[str(m)][j]=  matrix_mx_list[str(m)][j]+ matrix_mx_list[str(3)][j]
                            intervened_times_list[str(m)][j]= intervened_times_list[str(m)][j]+ intervened_times_list[str(3)][j]

                        if(j>2):
                            matrix_mx_list[str(m)][j]=  matrix_mx_list[str(0)][j]
                            intervened_times_list[str(m)][j]= matrix_mx_list[str(0)][j]

                    
                
                if (m==3):
                    for j in range(graph.num_parents_y):
                        if(j<3):
                            matrix_mx_list[str(m)][j]=  matrix_mx_list[str(2)][j]
                            intervened_times_list[str(m)][j]= intervened_times_list[str(2)][j]
                        if(j>2):
                            matrix_mx_list[str(m)][j]= matrix_mx_list[str(1)][j]
                            intervened_times_list[str(m)][j]= intervened_times_list[str(1)][j]

            if i % 1000 == 99:
                print("cluster")
              
                print("round:",i)

                print("reward:", total_expected_payoff)
                payoff_list.append(total_expected_payoff)

                
            # print(best_intervention)
        # print(np.matmul(np.linalg.inv(matrix_m), np.array([by]).T))
        return payoff_list, total_expected_payoff
   
def run_blm_lr_Fed_True(graph_list,best_y_list):
    graph=graph_list[0]
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}


    if graph.type == 'parallel':
        for m in range(len(graph_list)):
            client_num=len(graph_list)
            matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
            for i in range(graph.num_parents_y):
                matrix_m[i][i] = 1

            matrix_m_sum = np.zeros((graph.num_parents_y, graph.num_parents_y))
            for i in range(graph.num_parents_y):
                matrix_m_sum[i][i] = 1
            matrix_mx_sum = np.zeros(graph.num_parents_y)
            intervened_times_sum = np.zeros(graph.num_parents_y)
            by_sum = np.array([np.zeros(graph.num_parents_y)]).T
           
            matrix_m_list[str(m)]=  matrix_m
            matrix_mx = np.zeros(graph.num_parents_y)
            matrix_mx_list[str(m)]=  matrix_mx
      
        
            intervened_times = np.zeros(graph.num_parents_y)
            intervened_times_list[str(m)]=   intervened_times
            by = np.array([np.zeros(graph.num_parents_y)]).T
            by_list[str(m)]= by
            payoff_list = [0]
            total_expected_payoff = 0
            total_optimal_reward=0
        
        # we have a graph list
        #
        for i in range(graph.T):
            for m in range(len(graph_list)):
                matrix_m_list[str(m)]= matrix_m_sum
                matrix_mx_list[str(m)]=matrix_mx_sum
                intervened_times_list[str(m)]=intervened_times_sum
                by_list[str(m)] =by_sum

            for m in range(len(graph_list)):
                graph=graph_list[m]
                matrix_m=  matrix_m_list[str(m)]
                matrix_mx=  matrix_mx_list[str(m)]

                intervened_times=intervened_times_list[str(m)]
                by= by_list[str(m)] 



                inverse_m = np.linalg.inv(matrix_m)
                hat_theta_x = np.zeros(graph.num_parents_y)
                for j in range(graph.num_parents_y):
                    hat_theta_x[j] = matrix_mx[j] / (i- intervened_times[j]) if intervened_times[j] < i*client_num else 0
                hat_theta_x = np.array([hat_theta_x]).T
                
                norm_matrix_mx = np.array([np.sqrt(i*client_num - intervened_times[j]) for j in range(graph.num_parents_y)])
                hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
                rho = graph.compute_rho_lr(i+ 1)
                best_intervention = find_best_intervention(graph, inverse_m, hat_theta_x, hat_theta_y, norm_matrix_mx, rho)
                vy, y = graph.simulate(best_intervention)
                total_expected_payoff += graph.expect_y(best_intervention)
                total_optimal_reward+=best_y_list[m]
                
               
               
                matrix_m_sum += np.matmul(vy, vy.T)
                by_sum += y * vy
                for j in best_intervention:
                    intervened_times_sum[int(j)] += 1
                for j in range(graph.num_parents_y):
                    if j not in best_intervention:
                        matrix_mx_sum[j] += vy[j]

            

              
            if i % 100 == 99:
                print("True_Fed")
              
                print("round:",i)

                print("reward:", total_expected_payoff)
                total_regret=    total_optimal_reward-total_expected_payoff
                print("regret:",total_regret)
                payoff_list.append(total_regret)

                
            # print(best_intervention)
        # print(np.matmul(np.linalg.inv(matrix_m), np.array([by]).T))
        return payoff_list, total_expected_payoff

def run_blm_lr_Ind(graph_list,best_y_list):
    graph=graph_list[0]
    matrix_m_list={}
    matrix_mx_list={}
    intervened_times_list={}
    by_list={}

    if graph.type == 'parallel':
        for m in range(len(graph_list)):
            matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
            for i in range(graph.num_parents_y):
                matrix_m[i][i] = lambda_matrix
           
            matrix_m_list[str(m)]=  matrix_m
            matrix_mx = np.zeros(graph.num_parents_y)
            matrix_mx_list[str(m)]=  matrix_mx
      
        
            intervened_times = np.zeros(graph.num_parents_y)
            intervened_times_list[str(m)]=   intervened_times
            by = np.array([np.zeros(graph.num_parents_y)]).T
            by_list[str(m)]= by
            payoff_list = [0]
            total_expected_payoff = 0
            total_optimal_reward=0
        
        # we have a graph list
        #
        for i in range(graph.T):
            for m in range(len(graph_list)):
                
                graph=graph_list[m]
                matrix_m=  matrix_m_list[str(m)]
                matrix_mx=  matrix_mx_list[str(m)]

                intervened_times=intervened_times_list[str(m)]
                by= by_list[str(m)] 



                inverse_m = np.linalg.inv(matrix_m)
                hat_theta_x = np.zeros(graph.num_parents_y)
                for j in range(graph.num_parents_y):
                    hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
                hat_theta_x = np.array([hat_theta_x]).T
                norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])
                hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
                rho = graph.compute_rho_lr(i + 1)
                best_intervention = find_best_intervention(graph, inverse_m, hat_theta_x, hat_theta_y, norm_matrix_mx, rho)
                vy, y = graph.simulate(best_intervention)
                total_expected_payoff += graph.expect_y(best_intervention)
                total_optimal_reward+=best_y_list[m]
                
                
               
               
                matrix_m += np.matmul(vy, vy.T)
                by += y * vy
                for j in best_intervention:
                    intervened_times[int(j)] += 1
                for j in range(graph.num_parents_y):
                    if j not in best_intervention:
                        matrix_mx[j] += vy[j]

                matrix_m_list[str(m)]= matrix_m
                matrix_mx_list[str(m)]=matrix_mx
                intervened_times_list[str(m)]=intervened_times
                by_list[str(m)] =by
            if i % 100 == 99:
                print("Ind")
              
                print("round:",i)

                print("reward:", total_expected_payoff)

                total_regret=total_optimal_reward-total_expected_payoff

                print("regret:",   total_regret)
                payoff_list.append(   total_regret)

                
            # print(best_intervention)
        # print(np.matmul(np.linalg.inv(matrix_m), np.array([by]).T))
        return payoff_list, total_expected_payoff
    elif graph.type == 'two layer':
        matrix_m1 = np.zeros((3, 3))
        matrix_m2 = np.zeros((3, 3))
        matrix_m3 = np.zeros((3, 3))
        matrix_my = np.zeros((3, 3))
        for i in range(3):
            matrix_m1[i][i] = 1
            matrix_m2[i][i] = 1
            matrix_m3[i][i] = 1
            matrix_my[i][i] = 1
        matrix_mx = np.zeros(2)
        intervened_times = np.zeros(2)
        b1 = np.array([np.zeros(3)]).T
        b2 = np.array([np.zeros(3)]).T
        b3 = np.array([np.zeros(3)]).T
        by = np.array([np.zeros(3)]).T
        payoff_list = [0]
        total_expected_payoff = 0
        for i in range(graph.T):
            inverse_m1 = np.linalg.inv(matrix_m1)
            inverse_m2 = np.linalg.inv(matrix_m2)
            inverse_m3 = np.linalg.inv(matrix_m3)
            inverse_my = np.linalg.inv(matrix_my)
            hat_theta_1 = np.zeros(2)
            for j in range(2):
                hat_theta_1[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
            hat_theta_1 = np.array([hat_theta_1]).T
            norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(2)])
            hat_theta_2 = np.matmul(inverse_m1, np.array([b1]).T)
            hat_theta_3 = np.matmul(inverse_m2, np.array([b2]).T)
            hat_theta_4 = np.matmul(inverse_m3, np.array([b3]).T)
            hat_theta_y = np.matmul(inverse_my, np.array([by]).T)
            rho = graph.compute_rho_lr(i + 1)
            best_intervention = find_best_intervention_two_layer(graph, inverse_m1, inverse_m2, inverse_m3, inverse_my,
                                                                 hat_theta_1, hat_theta_2, hat_theta_3, hat_theta_4,
                                                                 hat_theta_y, norm_matrix_mx, rho)
            vy, y = graph.simulate(best_intervention)
            total_expected_payoff += graph.expect_y(best_intervention)
            if i % 100 == 99:
                payoff_list.append(total_expected_payoff)
            if 0 not in best_intervention:
                matrix_mx[0] += vy[0, 0]
            else:
                intervened_times[0] += 1
            if 1 not in best_intervention:
                matrix_mx[1] += vy[1, 0]
            else:
                intervened_times[1] += 1
            if 2 not in best_intervention:
                matrix_m1 += np.matmul(np.array([[1, vy[0, 0], vy[1, 0]]]).T, np.array([[1, vy[0, 0], vy[1, 0]]]))
                b1 += vy[2, 0] * np.array([[1, vy[0, 0], vy[1, 0]]]).T
            if 3 not in best_intervention:
                matrix_m2 += np.matmul(np.array([[1, vy[0, 0], vy[1, 0]]]).T, np.array([[1, vy[0, 0], vy[1, 0]]]))
                b2 += vy[3, 0] * np.array([[1, vy[0, 0], vy[1, 0]]]).T
            if 4 not in best_intervention:
                matrix_m3 += np.matmul(np.array([[1, vy[0, 0], vy[1, 0]]]).T, np.array([[1, vy[0, 0], vy[1, 0]]]))
                b3 += vy[4, 0] * np.array([[1, vy[0, 0], vy[1, 0]]]).T
            matrix_my += np.matmul(np.array([[vy[2, 0], vy[3, 0], vy[4, 0]]]).T,
                                   np.array([[vy[2, 0], vy[3, 0], vy[4, 0]]]))
            by += y * np.array([[vy[2, 0], vy[3, 0], vy[4, 0]]]).T
        return payoff_list, total_expected_payoff

def run_blm_lr_Fed(graph_list,best_y_list):
    graph=graph_list[0]
  

    if graph.type == 'parallel':
        matrix_m = np.zeros((graph.num_parents_y, graph.num_parents_y))
        for i in range(graph.num_parents_y):
            matrix_m[i][i] = 1
    
        matrix_mx = np.zeros(graph.num_parents_y)
    
    
        intervened_times = np.zeros(graph.num_parents_y)

        by = np.array([np.zeros(graph.num_parents_y)]).T

        payoff_list = [0]
        total_expected_payoff = 0
        total_optimal_reward=0
           
        
        # we have a graph list
        #
        for i in range(graph.T):
            for m in range(len(graph_list)):
                graph=graph_list[m]
           


                inverse_m = np.linalg.inv(matrix_m)
                hat_theta_x = np.zeros(graph.num_parents_y)
                for j in range(graph.num_parents_y):
                    hat_theta_x[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
                hat_theta_x = np.array([hat_theta_x]).T
                norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(graph.num_parents_y)])
                hat_theta_y = np.matmul(inverse_m, np.array([by]).T)
                rho = graph.compute_rho_lr(i + 1)
                best_intervention = find_best_intervention(graph, inverse_m, hat_theta_x, hat_theta_y, norm_matrix_mx, rho)
                vy, y = graph.simulate(best_intervention)
                total_expected_payoff += graph.expect_y(best_intervention)
                total_optimal_reward+=best_y_list[m]
                
               
                
                matrix_m += np.matmul(vy, vy.T)
                by += y * vy
                for j in best_intervention:
                    intervened_times[int(j)] += 1
                for j in range(graph.num_parents_y):
                    if j not in best_intervention:
                        matrix_mx[j] += vy[j]
            if i % 100 == 99:
                print("Fed")
                print("round:",i)
                print("reward:", total_expected_payoff)
                total_regret=     total_optimal_reward-total_expected_payoff
                payoff_list.append(   total_regret)

             

                
            # print(best_intervention)
        # print(np.matmul(np.linalg.inv(matrix_m), np.array([by]).T))
        return payoff_list, total_expected_payoff
    elif graph.type == 'two layer':
        matrix_m1 = np.zeros((3, 3))
        matrix_m2 = np.zeros((3, 3))
        matrix_m3 = np.zeros((3, 3))
        matrix_my = np.zeros((3, 3))
        for i in range(3):
            matrix_m1[i][i] = 1
            matrix_m2[i][i] = 1
            matrix_m3[i][i] = 1
            matrix_my[i][i] = 1
        matrix_mx = np.zeros(2)
        intervened_times = np.zeros(2)
        b1 = np.array([np.zeros(3)]).T
        b2 = np.array([np.zeros(3)]).T
        b3 = np.array([np.zeros(3)]).T
        by = np.array([np.zeros(3)]).T
        payoff_list = [0]
        total_expected_payoff = 0
        for i in range(graph.T):
            inverse_m1 = np.linalg.inv(matrix_m1)
            inverse_m2 = np.linalg.inv(matrix_m2)
            inverse_m3 = np.linalg.inv(matrix_m3)
            inverse_my = np.linalg.inv(matrix_my)
            hat_theta_1 = np.zeros(2)
            for j in range(2):
                hat_theta_1[j] = matrix_mx[j] / (i - intervened_times[j]) if intervened_times[j] < i else 0
            hat_theta_1 = np.array([hat_theta_1]).T
            norm_matrix_mx = np.array([np.sqrt(i - intervened_times[j]) for j in range(2)])
            hat_theta_2 = np.matmul(inverse_m1, np.array([b1]).T)
            hat_theta_3 = np.matmul(inverse_m2, np.array([b2]).T)
            hat_theta_4 = np.matmul(inverse_m3, np.array([b3]).T)
            hat_theta_y = np.matmul(inverse_my, np.array([by]).T)
            rho = graph.compute_rho_lr(i + 1)
            best_intervention = find_best_intervention_two_layer(graph, inverse_m1, inverse_m2, inverse_m3, inverse_my,
                                                                 hat_theta_1, hat_theta_2, hat_theta_3, hat_theta_4,
                                                                 hat_theta_y, norm_matrix_mx, rho)
            vy, y = graph.simulate(best_intervention)
            total_expected_payoff += graph.expect_y(best_intervention)
            if i % 100 == 99:
                payoff_list.append(total_expected_payoff)
            if 0 not in best_intervention:
                matrix_mx[0] += vy[0, 0]
            else:
                intervened_times[0] += 1
            if 1 not in best_intervention:
                matrix_mx[1] += vy[1, 0]
            else:
                intervened_times[1] += 1
            if 2 not in best_intervention:
                matrix_m1 += np.matmul(np.array([[1, vy[0, 0], vy[1, 0]]]).T, np.array([[1, vy[0, 0], vy[1, 0]]]))
                b1 += vy[2, 0] * np.array([[1, vy[0, 0], vy[1, 0]]]).T
            if 3 not in best_intervention:
                matrix_m2 += np.matmul(np.array([[1, vy[0, 0], vy[1, 0]]]).T, np.array([[1, vy[0, 0], vy[1, 0]]]))
                b2 += vy[3, 0] * np.array([[1, vy[0, 0], vy[1, 0]]]).T
            if 4 not in best_intervention:
                matrix_m3 += np.matmul(np.array([[1, vy[0, 0], vy[1, 0]]]).T, np.array([[1, vy[0, 0], vy[1, 0]]]))
                b3 += vy[4, 0] * np.array([[1, vy[0, 0], vy[1, 0]]]).T
            matrix_my += np.matmul(np.array([[vy[2, 0], vy[3, 0], vy[4, 0]]]).T,
                                   np.array([[vy[2, 0], vy[3, 0], vy[4, 0]]]))
            by += y * np.array([[vy[2, 0], vy[3, 0], vy[4, 0]]]).T
        return payoff_list, total_expected_payoff

if __name__ == "__main__":
    from two_layer_graph import TwoLayerGraph

    newGraph = TwoLayerGraph(10000)
    _, payoff = run_blm_lr(newGraph)
    print("_:",_)
    print(newGraph.best_expect_y() * 10000 - payoff)
